import librosa.display
import os
from moviepy.editor import VideoFileClip


def mp4_to_mp3(musicDir, input_mp4_file, output_mp3_file):
    if os.path.exists(output_mp3_file):
        print(f"{output_mp3_file} exists.")
        return
    if not os.path.exists(musicDir):
        os.mkdir(musicDir)
    # 打开MP4文件
    video_clip = VideoFileClip(input_mp4_file)
    # 提取音频
    audio_clip = video_clip.audio

    # 保存为MP3文件
    audio_clip.write_audiofile(output_mp3_file)
    # 关闭文件
    audio_clip.close()
    video_clip.close()

    print(f"MP3文件已生成: {output_mp3_file}")


def music_beat_time(output_mp3_file):
    # 提取音频
    y, sr = librosa.load(output_mp3_file)
    tempo, beat_frames = librosa.beat.beat_track(y=y, sr=sr)
    # 将卡点时间转换为时间（秒）
    beat_times = librosa.frames_to_time(beat_frames, sr=sr)
    results = []
    for beat_time in beat_times:
        results.append(str(round(beat_time, 2)))
        print(f'卡点时间：{beat_time:.2f} 秒')
    duration = librosa.get_duration(y=y, sr=sr)
    print(f'结束时间：{duration:.2f} 秒')
    results.append(str(round(duration, 2)))
    return results


if __name__ == '__main__':
    fileName = "video/UID1104351167776035_叫我小小张/video/2024-02-01 12.31.34-叫我小小张-人为什么总是喜欢反季节的事物# 甜妹 # 初恋"
    input_mp4_file = f"G:/{fileName}.mp4"
    fileName = os.path.basename(input_mp4_file).split(".mp4")[0]

    # fileName = "2023-12-30 21.15.31-呼呼-0326 j@c.nq 扣6吸走前任财运# 御姐"
    # UIDAuthor = "UID3510099090342622_呼呼"
    # # 输入的MP4文件名
    # input_mp4_file = f"E:/video/{UIDAuthor}/video/{fileName}.mp4"

    # 输出的MP3文件名
    musicDir = 'music'
    output_mp3_file = f"{musicDir}/{fileName}.mp3"
    beat_dir = 'beat'
    mp4_to_mp3(musicDir, input_mp4_file, output_mp3_file)

    # 提取音频数组
    beat_times = music_beat_time(output_mp3_file)
    # 打开文件并写入数组元素
    if not os.path.exists(beat_dir):
        os.mkdir(beat_dir)
    file_name_with_extension = os.path.basename(output_mp3_file)
    file_name, file_extension = os.path.splitext(file_name_with_extension)
    file_path = os.path.join(beat_dir, file_name + ".txt")
    with open(file_path, 'w', encoding="UTF-8") as file:
        for beat_time in beat_times:
            file.write(beat_time + '\n')
    os.system(f'notepad.exe {file_path}')
    print(fileName)
